#include "cube_signature.hpp"
#include "mpi/partition_3d.hpp"

static inline long next_powerof2(long x) {
  int log2x = std::floor(std::log(x) / std::log(2));
  long next = 1l << log2x;
  return next < x ? next * 2l : next;
}

cube_signature get_expected_mnk(long m, long n, long k) {
  bool is_square = m == n && m == k;
  long hyper_m = is_square ? m : next_powerof2(m);
  long hyper_n = is_square ? n : next_powerof2(n);
  long hyper_k = is_square ? k : next_powerof2(k);
  return std::move(cube_signature({hyper_m, hyper_n, hyper_k}));
}

void argsParser(int argc, char *argv[], int &p, int &M, int &N, int &K, int &mp,
                int &np, int &kp) {
  if (argc < 6) {
    INFO_PRINTF("argc = %d\n", argc);
    ERROR_PRINTF("Usage: %s <p> <M> <N> <K> <OPT>\n", argv[0]);
    exit(1);
  }

  p = std::atoi(argv[1]);
  M = std::atoi(argv[2]);
  N = std::atoi(argv[3]);
  K = std::atoi(argv[4]);
  int opt = std::atoi(argv[5]);

  const int opt_sq = opt * opt;
  long bvm, bvn, bvk;
  if (K >= opt) {
    if (M >= opt && N >= opt) {
      bvm = M / opt;
      bvn = N / opt;
      bvk = K / opt;
    } else if (M < opt && N < opt) {
      bvm = 1;
      bvn = 1;
      bvk = K / opt;
    } else if (M < opt || N < opt) {
      if (M < opt) {
        bvm = 1;
        bvn = M * N / opt_sq;
        bvk = K / M;
      } else {
        bvn = 1;
        bvm = M * N / opt_sq;
        bvk = K / N;
      }
    }
  } else {
    int *x, *y, *z;
    long *bvx, *bvy, *bvz;
    if (M >= N && N >= K) {
      x = &M, y = &N, z = &K, bvx = &bvm, bvy = &bvn, bvz = &bvk;
    } else if (M >= K && K >= N) {
      x = &M, y = &K, z = &N, bvx = &bvm, bvy = &bvk, bvz = &bvn;
    } else if (N >= M && M >= K) {
      x = &N, y = &M, z = &K, bvx = &bvn, bvy = &bvm, bvz = &bvk;
    } else if (N >= K && K >= M) {
      x = &N, y = &K, z = &M, bvx = &bvn, bvy = &bvk, bvz = &bvm;
    } else {
      ERROR_PRINTF("Single CGEMM is enough\n");
    }
    int P1 = (*x) * (*y) / opt_sq, P2 = (*y) / (*z);
    if (M >= opt && N >= opt) {
      *bvy = 1 << static_cast<long>(std::floor(std::log2(M * N / opt_sq) / 2));
      *bvx = M * N / opt_sq / (*bvy);
      *bvz = 1;
    } else {
      *bvy = std::min(P1, P2);
      if (*bvy == P1) {
        *bvy = 1 << static_cast<long>(std::floor(std::log2(P1) / 2));
      }
      *bvx = P1 / (*bvy), *bvz = 1;
    }
  }
  auto eh_mnk = get_expected_mnk(M / bvm, N / bvn, K / bvk);
  int R = std::min(eh_mnk.size(0), std::min(eh_mnk.size(1), eh_mnk.size(2)));
  Solution optimal = findOptimal(p, bvm, bvn, bvk, R, true);
  mp = optimal.pm, np = optimal.pn, kp = optimal.pk;
  if ((mp < 1) || (np < 1) || (kp < 1) || (mp * np * kp > p)) {
    ERROR_PRINTF(
        "Invalid process grid generated: p = %d, mp, np, kp = %d, %d, %d\n", p,
        mp, np, kp);
    exit(1);
  }

  printf("Problem Size m   * n   * k   :\t%-5d * %-5d * %-5d\n", M, N, K);
  printf("Block View   bm  * bn  * bk  :\t%-5ld * %-5ld * %-5ld\n", bvm, bvn,
         bvk);
  printf("BV Cuboid    bbm * bbn * bbk :\t%-5ld * %-5ld * %-5ld\n",
         eh_mnk.size(0), eh_mnk.size(1), eh_mnk.size(2));
  printf("Process Grid mp  * np  * kp  :\t%-5d * %-5d * %-5d\n", mp, np, kp);
  if (M % mp != 0 || K % kp != 0 || N % np != 0) {
    ERROR_PRINTF("Invalid problem size: mp, np, kp = %d, %d, %d must excatly "
                 "divide M, K and N\n",
                 mp, np, kp);
    exit(1);
  }
  int mb = M / mp;
  int nb = N / np;
  int kb = K / kp;
  printf("Work cuboid  mb  * nb  * kb  :\t%-5d * %-5d * %-5d\n", mb, nb, kb);
  fflush(stdout);
}

int main(int argc, char *argv[]) {
  // Get MPI metadata
  int p;

  int M, K, N;
  int mp, kp, np;
  argsParser(argc, argv, p, M, N, K, mp, np, kp);
  return 0;
}